# -*- coding: utf-8 -*-
import os

import psycopg2
import pymysql

current_path = os.path.dirname(__file__)
os.chdir(current_path)


def create_data_list(data_root_path):
    with open(os.path.join(data_root_path, 'dict_txt.txt'), 'r', encoding='utf-8') as f_data:
        dict_txt = eval(f_data.readlines()[0])

    lines = getDataByPostGres("qa_train")
    # lines = getDataBysql("qa_train")
    i = 0
    for line in lines:
        title = line[2]
        l = line[1]
        labs = ""
        with open(os.path.join(data_root_path, 'train_list.txt'), 'a', encoding='utf-8') as f_train:
            for s in title:
                if s != " ":
                    lab = str(dict_txt[s])
                    labs = labs + lab + ','
            labs = labs[:-1]
            labs = labs + '\t' + str(l) + '\n'
            f_train.write(labs)
        i += 1
    print("数据列表生成完成！")


def getDataBysql(tableName):
    connect = pymysql.Connect(
        host="192.168.8.23",
        port=3306,
        user="root",
        passwd="root",
        db="ld_qa",
        charset='utf8'
    )
    cursor = connect.cursor()
    sql = "select * from " + tableName
    cursor.execute(sql)
    lines = cursor.fetchall()
    cursor.close()
    connect.close()
    return lines


def getDataByPostGres(tableName):
    connect = psycopg2.connect(
        database="jjgj_db_23", user="postgres", password="123456", host="192.168.8.25", port="5432"
    )
    cursor = connect.cursor()
    sql = "select * from " + tableName
    cursor.execute(sql)
    lines = cursor.fetchall()
    cursor.close()
    connect.close()
    return lines


def create_dict(dict_path):
    dict_set = set()
    lines = getDataByPostGres("qa_train")
    # lines = getDataBysql("qa_train")
    for line in lines:
        title = line[2]
        for s in title:
            if s != " ":
                dict_set.add(s)
    # 把元组转换成字典，一个字对应一个数字
    dict_list = []
    i = 0
    for s in dict_set:
        dict_list.append([s, i])
        i += 1
    # 添加未知字符
    dict_txt = dict(dict_list)
    end_dict = {"<unk>": i}
    dict_txt.update(end_dict)
    # 把这些字典保存到本地
    with open(dict_path, 'w', encoding='utf-8') as f:
        f.write(str(dict_txt))
    print("数据字典生成成功")


data_root_path = "../dataset"


def start_prepare_dataset():
    if os.path.exists(data_root_path + '/dict_txt.txt'):  # 如果文件存在
        os.remove(data_root_path + '/dict_txt.txt')
        os.remove(data_root_path + '/train_list.txt')
    else:
        print('no such file:', "dict_txt.txt,train_list.txt")
    dict_path = os.path.join(data_root_path, 'dict_txt.txt')
    # 创建数据字典
    create_dict(dict_path)
    # 创建数据列表
    create_data_list(data_root_path)


if __name__ == '__main__':
    start_prepare_dataset()
    print('数据部署成功')
